// Copyright (c) 2025 Yao Yu
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use alloc::{collections::BinaryHeap, boxed::Box, sync::Arc};

use core::cmp::{self, Reverse};

use ostd::{
    arch::read_tsc,
    cpu::CpuId,
    sync::SpinLock,
    task::{
        scheduler::{
            LocalRunQueue,
            EnqueueFlags, Scheduler, inject_scheduler,
            UpdateFlags,
        },
        Task,
    },
    trap::disable_local,
};

pub fn sched_init() {
    let scheduler: &'static mut dyn Scheduler = Box::leak(Box::new(ClassScheduler::new()));
    inject_scheduler(scheduler);
    // TODO: support resource load statistic
}

// 带有任务id的一个任务项。这个id暂时用于小根堆的排序key
struct FairQueueItem(Arc<Task>, u64);
// 下面为FairQueueItem实现了PartialOrd、Ord、Debug、PartialEq、Eq等trait，方便排序和比较。
impl core::fmt::Debug for FairQueueItem {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{:?}", self.key())
    }
}

impl FairQueueItem {
    fn key(&self) -> u64 {
        self.1
    }
}

impl PartialEq for FairQueueItem {
    fn eq(&self, other: &Self) -> bool {
        self.key().eq(&other.key())
    }
}

impl Eq for FairQueueItem {}

impl PartialOrd for FairQueueItem {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for FairQueueItem {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.key().cmp(&other.key())
    }
}

// 调度器结构体，将会为之实现ostd::task::scheduler::Scheduler trait
pub struct ClassScheduler {
    // 似乎每个cpu核心都有自己的队列？
    rqs: Box<[SpinLock<RunQueue>]>,
}

struct CurrentRuntime {
    start: u64,         // 时间片开始时间
    period_delta: u64,  // 从start到现在，运行了多少时间（每次更新的时候，这个东西加一点）
}
impl CurrentRuntime {
    fn new() -> Self {
        CurrentRuntime {
            start: read_tsc(),
            period_delta: 0,
        }
    }
    fn update(&mut self) {
        let new_now = read_tsc();
        let delta = new_now - core::mem::replace(&mut self.start, new_now);
        self.period_delta += delta;
    }
}
const TIME_SLICE: u64 = 100;

// 任务队列，以及一些metadata
pub struct RunQueue {
    entities: BinaryHeap<Reverse<FairQueueItem>>,   // 任务队列(待运行)
    current: Option<(Arc<Task>, CurrentRuntime)>,        // 当前运行的任务
    taskid: u64,                                    // 任务id
}

impl RunQueue {
    fn new() -> Self {
        Self {
            entities: BinaryHeap::new(),
            current: None,
            taskid: 0,
        }
    }

    // 向任务队列中添加任务
    fn enqueue_entities(&mut self, task: Arc<Task>) {
        let ritem = Reverse(FairQueueItem(task, self.taskid));
        self.entities.push(ritem);
        self.taskid+=1;
    }

    // 从队列中pop一个拿出来
    fn pick_next(&mut self) -> Option<Arc<Task>> {
        let Reverse(FairQueueItem(task, _)) = self.entities.pop()?;
        Some(task)
    }
}

// 任务队列必须实现的trait
impl LocalRunQueue for RunQueue {
    
    // 返回当前正在运行的task
    fn current(&self) -> Option<&Arc<Task>> {
        self.current
        .as_ref()
        .map(|(task, _)| task)
    }

    // 更新当前运行任务的调度统计信息（以及可能任务在队列中的位置）
    // 如果当前可运行的任务需要被抢占，则返回true
    // TODO: 更新调度统计信息，去掉flags前面的_
    fn update_current(&mut self, flags: UpdateFlags) -> bool {
        if let Some((_, crt)) = &mut self.current {
            crt.update();
            // 主动释放应该被抢占，超过时间片应该被抢占
            match flags {
                UpdateFlags::Yield => true,
                UpdateFlags::Tick | UpdateFlags::Wait => {
                    crt.period_delta >= TIME_SLICE
                }
            }
        } else {
            true
        }
    }

    // 任务队列中pop一个新的任务出来
    // TODO: 旧的任务可以考虑先放在队列尾部
    fn pick_next_current(&mut self) -> Option<&Arc<Task>> {
        // 把新的任务拿出来，放到current里面
        self.pick_next().and_then(|next| {
            if let Some(old) = self.current.replace((next, CurrentRuntime::new())){
                // 把旧的current任务放回去
                self.enqueue_entities(old.0);
            }
            // 返回当前任务
            self.current()
        })
    }

    // 似乎只需要将当前任务的cpu设置为None并将任务返回即可？
    fn dequeue_current(&mut self) -> Option<Arc<Task>> {
        self.current.take().map(|(cur_task, _)| {
            cur_task.schedule_info().cpu.set_to_none();
            cur_task
        })
    }
}

impl ClassScheduler {
    pub fn new() -> Self {
        ClassScheduler {
            rqs: Box::new([SpinLock::new(RunQueue::new())]),
        }
    }
}

impl Scheduler for ClassScheduler {
    fn enqueue(&self, runnable: Arc<Task>, _flags: EnqueueFlags) -> Option<CpuId> {
        // TODO: support more cpus
        // TODO: support flags
        let current_cpu = CpuId::bsp();
        // hint: 这里一定要disable_irq，否则这里如果占据任务队列的锁后，触发时钟中断后再次请求队列锁，会导致死锁
        self.rqs[current_cpu.as_usize()].disable_irq().lock().enqueue_entities(runnable);
        Some(CpuId::bsp())
    }

    fn local_rq_with(&self, f: &mut dyn FnMut(&dyn LocalRunQueue)) {
        // TODO: support more cpus
        // TODO: guard in asterinas:asterinas/kernel/src/sched/sched_class/mod.rs:231, 需要了解为什么disable_local
        let _guard = disable_local();
        let current_cpu = CpuId::bsp();
        f(&*self.rqs[current_cpu.as_usize()].lock())
    }

    fn local_mut_rq_with(&self, f: &mut dyn FnMut(&mut dyn LocalRunQueue)) {
        // TODO: support more cpus
        // TODO: guard in asterinas:asterinas/kernel/src/sched/sched_class/mod.rs:231, 需要了解为什么disable_local
        let _guard = disable_local();
        let current_cpu = CpuId::bsp();
        let mut lock = self.rqs[current_cpu.as_usize()].lock();
        f(&mut *lock)
    }
}